import torch
from model.Unet_model import Unet_model
from model.utils import DownSampling, ConvBlock, UpSampling
from args import get_args
import torch.nn as nn
from dataloader.dataload import getDataloader
from tqdm import tqdm
import numpy as np
from loss.loss import dice_loss,MIOU
import torch.nn.functional as F
import logging

logging.basicConfig(level=logging.INFO)


def test(args):
    test_loader = getDataloader(root=args.data_path+'/data',
                                txtfile_path=args.data_path + '/test.txt',
                                BatchSize=args.batch_size,
                                shuffle=False,
                                num_workers=args.num_workers)

    model = Unet_model(ConvBlock, DownSampling, UpSampling).to(args.device)
    model.load_state_dict(torch.load(args.ckpt_path))
    creterion = nn.CrossEntropyLoss(reduction="mean")
    # 训练过程
    logging.info(f"""start to test:
                     epoches:{args.epoches}
                     batch_size:{args.batch_size}
                     lr={args.lr}
                     num_workers={args.num_workers}
                     device={args.device}
                 """)
    with torch.no_grad():
        model.eval()
        with tqdm(total=len(test_loader), mininterval=1)as pbar:
            miou_list=[]
            test_loss_list = []
            for i, (img, mask) in enumerate(test_loader):
                img = img.to(args.device)
                mask = mask.to(args.device)
                output = model(img)
                loss = creterion(output, mask) + dice_loss(F.softmax(output, dim=1).float(),
                                                           F.one_hot(mask, 3).permute(0, 3, 1, 2).float(),
                                                           multiclass=True)
                miou=MIOU(F.softmax(output, dim=1).float(),F.one_hot(mask, 3).permute(0, 3, 1, 2).float())
                miou_list.append(miou.cpu().data)
                test_loss_list.append(loss.cpu().data)
                pbar.update()
            test_mean_loss = np.array(test_loss_list).mean()
            test_miou=np.array(miou_list).mean()
        print(f"test_loss:{test_mean_loss:.4f}")
        print(f"test_miou:{test_miou:.4f}")


if __name__ == '__main__':
    args = get_args()
    test(args)
